import tensorflow as tf
import numpy as np
import pandas as pd
import pdb
import json
from mod_core_rnn_cell_impl import LSTMCell  # modified to allow initializing bias in lstm

import data_utils
import plotting
import model
import mmd
import utils
import eval
import DR_discriminator

from differential_privacy.dp_sgd.dp_optimizer import dp_optimizer
from differential_privacy.dp_sgd.dp_optimizer import sanitizer
from differential_privacy.privacy_accountant.tf import accountant

import AutoencoderFunctions
import os
from time import time

begin = time()

"""
Here, both the discriminator and generator were used to do the anomaly detection
"""

# --- get settings --- #
# parse command line arguments, or use defaults
parser = utils.rgan_options_parser()
settings = vars(parser.parse_args())
# if a settings file is specified, it overrides command line arguments/defaults
if settings['settings_file']: settings = utils.load_settings_from_file(settings)

# --- get data, split --- #
data_path = './experiments/data/' + settings['data_load_from'] + '.data.npy'
print('Loading data from', data_path)
samples, labels, index = data_utils.get_data(settings["data"], settings["seq_length"], settings["seq_step"],
                                             settings["num_signals"], settings["sub_id"], settings["eval_single"],
                                             settings["eval_an"], data_path)
# --- save settings, data --- #
# no need
print('Ready to run with settings:')
for (k, v) in settings.items(): print(v, '\t', k)
# add the settings to local environment
# WARNING: at this point a lot of variables appear
locals().update(settings)
json.dump(settings, open('./experiments/settings/' + identifier + '.txt', 'w'), indent=0)

class myADclass():
    def __init__(self, epoch, settings=settings, samples=samples, labels=labels, index=index):
        self.epoch = epoch
        self.settings = settings
        self.samples = samples
        self.labels = labels
        self.index = index
    def ADfunc(self):
        num_samples_t = self.samples.shape[0]
        # t_size = 500  Decomentar, esse é o certo para considerar todos os batches
        t_size = 500
        T_index = np.random.choice(num_samples_t, size=t_size, replace=False)
        print('sample_shape:', self.samples.shape[0])
        print('num_samples_t', num_samples_t)

        # -- only discriminate one batch for one time -- #
        D_test = np.empty([t_size, self.settings['seq_length'], 1])
        DL_test = np.empty([t_size, self.settings['seq_length'], 1])
        GG = np.empty([t_size, self.settings['seq_length'], self.settings['num_signals']])
        T_samples = np.empty([t_size, self.settings['seq_length'], self.settings['num_signals']])
        L_mb = np.empty([t_size, self.settings['seq_length'], 1])
        I_mb = np.empty([t_size, self.settings['seq_length'], 1])
 
        for batch_idx in range(0, t_size):
            print('epoch:{}'.format(self.epoch))
            print('batch_idx:{}'.format(batch_idx))
            # display batch progress
            model.display_batch_progression(batch_idx, t_size)
            T_mb = self.samples[T_index[batch_idx], :, :]
            L_mmb = self.labels[T_index[batch_idx], :, :]
            I_mmb = self.index[T_index[batch_idx], :, :]
            para_path = './experiments/parameters/' + self.settings['sub_id'] + '_' + str(
                self.settings['seq_length']) + '_' + str(self.epoch) + '.npy'
            D_t, L_t = DR_discriminator.dis_D_model(self.settings, T_mb, para_path)
            Gs, Zs, error_per_sample, heuristic_sigma = DR_discriminator.invert(self.settings, T_mb, para_path,
                                                                                g_tolerance=None,
                                                                                e_tolerance=0.1, n_iter=None,
                                                                                max_iter=10,
                                                                                heuristic_sigma=None)


            GG[batch_idx, :, :] = Gs
            T_samples[batch_idx, :, :] = T_mb
            D_test[batch_idx, :, :] = D_t
            DL_test[batch_idx, :, :] = L_t
            L_mb[batch_idx, :, :] = L_mmb
            I_mb[batch_idx, :, :] = I_mmb

        #------------------------------------------------------------
        savePathPart1 = "./Experiments_Autoencoder2/MADGANInvert_CombinedDG/DR_Losses_EpochGAN" + str(self.epoch) + "/"
        os.mkdir(savePathPart1)
        savePath_DL1 = savePathPart1 + "/DL1" + ".npy"
        savePath_DL2 = savePathPart1 + "/DL2" + ".npy"
        savePath_LL = savePathPart1 + "/LL" +  ".npy"
        savePath_RL = savePathPart1 + "/RL" +  ".npy"
        D_L_1, D_L_2, L_L, R_L = AutoencoderFunctions.computeAndSaveDandRLosses(D_test, DL_test, GG, T_samples, L_mb, self.settings['seq_step'], savePath_DL1, savePath_DL2, savePath_LL, savePath_RL)
        #------------------------------------------------------------


        return


if __name__ == "__main__":
    print('Main Starting...')

    # for epoch in range(settings['num_epochs']):
    # for epoch in range(100):
    # for epoch in [21, 22]:
    for epoch in [8]:
        print("epoch: ", epoch)
        ob = myADclass(epoch)
        ob.ADfunc()
        end = time() - begin
        print('Execution time for one epoch = %d s' % (end))

    print('Main Terminating...')
    end = time() - begin
    print('Testing terminated | Execution time=%d s' % (end))